import pandas as pd
from pandas import read_excel

from data.CountryAfrica import countryAfrica


class WorldBankExcelHandler:

    def __init__(self, input_path):
        self.input_path = input_path
        self.output_sheet_name = input_path.split('_')[-1].split('.')[0]
        output_filename = '非洲_' + self.output_sheet_name + '.xls'
        self.output_path = input_path.rpartition('\\')[0] + '\\' + output_filename
        self.data = None
        self.save_data = None

    def read_excel(self, sh_name='Data', skip_row=3):
        self.data = pd.read_excel(self.input_path, sheet_name=sh_name, skiprows=skip_row)
        # filtered_df = df.loc[df['Country Code'].isin(list_country)]
        return self.data

    def africa_filter(self):
        abbrs = countryAfrica.get_sorted_keys()
        self.save_data = self.data.loc[self.data['Country Code'].isin(abbrs)]
        return self.save_data

    def display_data(self):
        if self.data is not None:
            return self.data.head()
        else:
            return "No data loaded."

    def add_column_to_save_data(self, column_name, values):
        if self.save_data is not None:
            # self.save_data['Chinese Name'] = self.save_data['Country Code'].map(countryAfrica.countries_chn)
            self.save_data.insert(1, 'Chinese Name', self.save_data['Country Code'].map(countryAfrica.countries_chn))
            # self.save_data[column_name] = values
            return "Column added successfully."
        else:
            return "No data loaded."

    def save_to_excel(self):
        if self.save_data is not None:
            with pd.ExcelWriter(self.output_path, engine='openpyxl') as writer:
                self.save_data.to_excel(writer, index=False, sheet_name=self.output_sheet_name)

            # self.save_data.to_excel(output_file_path, index=False)
            print("Data saved successfully.")
        else:
            print("No data loaded.")

    def read_filter_save(self):
        self.read_excel(sh_name='Data', skip_row=3)
        self.africa_filter()
        self.add_column_to_save_data(None, None)
        self.save_to_excel()


# 测试对象
if __name__ == '__main__':
    handler = WorldBankExcelHandler(
        r'D:\01.DOC\漏损产品\市场分析\API_SP.POP.TOTL_DS2_en_excel_v2_9_总人口.xls')
    handler.read_filter_save()

    pass
